Implement Trie

Implement a trie with insert, search, and startsWith methods.

Note:

You may assume that all inputs are consist of lowercase letters a-z.

Solution:

  1. class TrieNode {
  2.   char c;
  3.   boolean isLeaf;
  4.   Map<Character, TrieNode> children = new HashMap<>();
  6.   public TrieNode() {}
  7.   public TrieNode(char c) {
  8.     this.c = c;
  9.   }
  10. }
  12. public class Trie {
  13.   private TrieNode root;
  15.   public Trie() {
  16.     root = new TrieNode();
  17.   }
  19.   // Inserts a word into the trie.
  20.   public void insert(String word) {
  21.     Map<Character, TrieNode> children = root.children;
  23.     for (int i = 0; i < word.length(); i++) {
  24.       char c = word.charAt(i);
  26.       if (!children.containsKey(c)) {
  27.         children.put(c, new TrieNode(c));
  28.       }
  30.       if (i == word.length() - 1) {
  31.         children.get(c).isLeaf = true;
  32.       }
  34.       children = children.get(c).children;
  35.     }
  36.   }
  38.   // Returns if the word is in the trie.
  39.   public boolean search(String word) {
  40.     TrieNode node = getLastTrieNode(word);
  41.     return node != null && node.isLeaf;
  42.   }
  44.   // Returns if there is any word in the trie
  45.   // that starts with the given prefix.
  46.   public boolean startsWith(String prefix) {
  47.     TrieNode node = getLastTrieNode(prefix);
  48.     return node != null;
  49.   }
  51.   private TrieNode getLastTrieNode(String word) {
  52.     Map<Character, TrieNode> children = root.children;
  54.     for (int i = 0; i < word.length(); i++) {
  55.       char c = word.charAt(i);
  57.       if (!children.containsKey(c)) {
  58.         break;
  59.       }
  61.       if (i == word.length() - 1)
  62.         return children.get(c);
  64.       children = children.get(c).children;
  65.     }
  67.     return null;
  68.   }
  69. }
  71. // Your Trie object will be instantiated and called as such:
  72. // Trie trie = new Trie();
  73. // trie.insert("somestring");
  74. // trie.search("key");